#-----------------------------------------------------------------------------
#
# State-level results for Callaway and Li (2023)
#
# This code reproduces Figure 1 in the main text and
# Figures S-3 and S-6 in the Supplementary Appendix.
# 
#
# Authors: Brantly Callaway and Tong Li
# Date:    2/7/2023
#
#-----------------------------------------------------------------------------

library(data.table)
library(fixest)
library(ggalt)
library(ggplot2)
library(ggrepel)

load("state_covid_data.RData")

start_date <- as.Date("2020-03-16")
end_date <- as.Date("2020-05-10")

team.colors <- c("#828A8f", "#9D2235", "#BA0C2F", "#0033A0", "#006BA6", "#7BAFD4", "#FF8200") 

make_state_plot <- function(var, ylabel, save_plot=FALSE) {
  # per capita tests
  p <- ggplot(state_covid_data, aes(x=date, y=.data[[var]], group=state, color=state)) +
    geom_line(linewidth=1.1) + 
    ylab(ylabel) +
    xlab("Date") +
    theme_bw() +
    theme(legend.position="none") +
    scale_color_manual(values=team.colors) + 
    geom_label_repel(aes(label = state.label),
                     nudge_x = 0,
                     na.rm = TRUE) +
    scale_x_date(breaks=seq(start_date+5,end_date,by=7), date_labels="%b %d", limits=c(start_date+1, end_date-1))
  print(p)
  if (save_plot) {
    ggsave(paste0(var,".pdf"), width=6, height=3.5)
  }
}

save_plot <- FALSE

# produces Figure 1, panel (a)
make_state_plot("cum_tests.pc", "Cumulative Tests per 1000 People", save_plot=save_plot)

# produces Figure 1, panel (b)
make_state_plot("tests.pc.ma", "7 Day Moving Avg. Tests per 1000 People", save_plot=save_plot)

# produces Figure 1, panel (c)
make_state_plot("cum_cases.pc", "Cumulative Confirmed Cases per 1000 People", save_plot=save_plot)

# produces Figure 1, panel (d)
make_state_plot("cases.pc.ma", "7 Day Moving Avg. Confirmed Cases per 1000 People", save_plot=save_plot)

# produces Figure S-3, panel (a)
make_state_plot("cum_deaths.pc", "Cumulative Deaths per 1000 People", save_plot=save_plot)

# produces Figure S-3, panel (b)
make_state_plot("deaths.pc.ma", "7 Day Moving Avg. Deaths per 1000 People", save_plot=save_plot)

# produces Figure S-3, panel (c)
make_state_plot("cum_hosp.pc", "Cumulative Hospitalizations per 1000 People", save_plot=save_plot)



#-----------------------------------------------------------------------------
# code to produce bounds in Figure S-6
#-----------------------------------------------------------------------------

this_date <- as.Date("2020-03-31")
#this_date <- as.Date("2020-04-25")

this_data.dta <- subset(state_covid_data, date==this_date)
this_data.dta <- subset(this_data.dta, cum_cases.pc > 0)

# terms that bounds depend on 
p.R <- this_data.dta$cum_cases.pc/1000
p.T <- this_data.dta$cum_tests.pc/1000
p.RcondT1 <- this_data.dta$cum_cases.pc/this_data.dta$cum_tests.pc

# false negative rate of test
false.neg <- 0.25

gam <- p.R/(1-false.neg)
p.CcondT1 <- p.RcondT1/(1-false.neg)

Lower <- gam
# bound without any assumptions
Upper <- gam + (1-p.T)

# bound under assumption 1 in paper
Upper2 <- gam + p.CcondT1*(1-p.T)


# store results in data frame for plotting
this_data.dta$Upper <- Upper2
this_data.dta$Lower <- Lower

# make plot alphabetical
this_data.dta$state <- droplevels(as.factor(this_data.dta$state))
olevels <- levels(this_data.dta$state)
this_data.dta$state <- factor(this_data.dta$state, levels=olevels[-order(olevels)+length(olevels)+1])


ggplot(this_data.dta, aes(x=Lower, xend=Upper, y=state, group=state)) +
  geom_dumbbell() +
  geom_point(aes(color=state), size=6) +
  geom_point(aes(x=Upper, color=state), size=6) +
  scale_color_manual(values=rev(team.colors)) + 
  xlab("Infection Rate") +
  ylab("State") +
  xlim(c(0,.4)) +
  theme_bw() +
  theme(legend.position="none")
